#!/usr/bin/python
'''
    Filename - troll.py
    Author: Gerrit Kitts <gkitts@gmail.com>        
    Usage: troll.py <listening port> <datagram error percentage>
        <listening port> must be a valid IP port from 1024-65525
        <error percentage> must be an integer representation of an error percentage
    
    This package provides a basic UDP proxy that listens for incoming datagrams and forwards
    them to their ultimate destination.  A datagram may be dropped (i.e. ignored) 
    if a random value is less than or equal to the configured error value.  
    Each UDP packet consists of:
        remote IP <4 bytes>: The destination IP for the datagram
        remote Port <2 bytes>: The destination port for the datagram
        datagram type <1 byte>: 1 if the datagram contains the file size, 2 if the datagram
            contains the file name, 3 for data segments.
'''

import random
import socket
import struct

from sys import argv

BUFF_SIZE = 1024

'''
    Exception sub-class to handle invalid argument specifications
'''
class ArgumentException(BaseException):
    def __init__(self):
        self.msg = "Usage: troll.py (listeningPort) (errorRate)"
        
'''
    Exception sub-class to handle invalid data transfers
'''
class ForwardException(BaseException):
    def __init__(self):
        self.msg = "Error forwarding datagram to server!"

'''
    Creates a socket on all IP addresses on the port specified as a parameter, and listens
    for incoming requests.  
'''
def listenForConnections(params):
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    address = params['port']
    print "Binding socket %s:%s" % (address[0], address[1])
    s.bind(address)
    
    print "Waiting for packets"
    while True:
        datagram = s.recv(BUFF_SIZE)
        print "Received a datagram"
        
        # When we receive a connection, we might drop a packet.  Test to see if that's the case.
        rint = random.randint(1,100)
        if (rint <= params['errorRate']):
            print "Dropping datagram due to error (rate = %i, int = %i)" \
                % (params['errorRate'], rint)
        else:
            # Crack the datagram, and forward the connection.  We're interested in the first 
            #  6 bytes (remote IP + port) at this point.
            unpacker = struct.Struct("! 4BH")
            metadata = unpacker.unpack(datagram[0:unpacker.size])
            rAddr = ('.'.join([str(x) for x in metadata[0:4]]), metadata[4])
            print "Remote destination is %s:%s" % (rAddr)
        
            # Now, create a socket to the remote host and send the datagram
            forwardSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            numSent = forwardSock.sendto(datagram, rAddr)
            print "Forwarded %i bytes." % numSent
            forwardSock.close()
        
            # Ensure the bytes sent equal the datagram size
            if (numSent != len(datagram)):
                raise ForwardException
    
    print "Closing troll."
    s.close()

'''
    Base wrapper for class functionality
'''
if __name__ == '__main__':
    params = {
        'port' : ('', 7775),
        'errorRate' : 0
    }

    if len(argv) is 2:
        params['port'] = ('', int(argv[1]))
    if len(argv) is 3:
        params['errorRate'] = int(argv[2])
        
    # Initialize the PRNG
    random.seed()
    
    listenForConnections(params)
    
